import torch
import torch.nn as nn
from tqdm import tqdm
import torch.nn.functional as F

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

def fit_one_epoch(model_train, model, loss, optimizer, epoch, epoch_step, epoch_step_val, gen, genval, Epoch, cuda):
    total_loss      = 0
    total_accuracy  = 0

    val_loss            = 0
    val_total_accuracy  = 0
    
    model_train.train()
    print('Start Train')
    #tqdm是一个迭代对象的进度条库
    with tqdm(total=epoch_step,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
        # with torch.no_grad():
        # p = len(gen)
        for iteration, [images, targets] in enumerate(gen):
            if iteration >= len(gen):
                break
            # images, targets = batch[0], batch[1]
            # images = F.interpolate(images, size=(224, 224), mode='nearest')
            with torch.no_grad():
                if cuda:
                    images  = images.type(torch.FloatTensor).cuda()
                    targets = torch.tensor(targets).cuda()
                else:
                    images  = images.type(torch.FloatTensor)
                    targets = torch.tensor(targets)

            images = F.interpolate(images, size=(224, 224), mode='nearest')

            # print(images)
            optimizer.zero_grad()#梯度置零

            #--------------------
            #这里的image为(2,8,3,720,1280),image[0]为left,image[1]为right,image[0][0]和image[1][0]为一对
            #-------------------

            # m = model_train(images)
            outputs = nn.Sigmoid()(model_train(images))
            output  = loss(outputs, targets)

            output.backward()
            optimizer.step()#更新参数

            with torch.no_grad():
                _, preds = torch.max(outputs, 1)
                equal       = torch.sum(preds.view(-1) == targets.view(-1))
                accuracy = equal / targets.__len__()

            total_loss      += output.item()
            total_accuracy  += accuracy.item()

            pbar.set_postfix(**{'total_loss': total_loss / (iteration + 1),
                                'acc'       : total_accuracy / (iteration + 1),
                                'lr'        : get_lr(optimizer)})
            pbar.update(1)
    print('Finish Train')

    model_train.eval()
    print('Start Validation')
    with tqdm(total=epoch_step_val, desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
        for iteration, batch in enumerate(genval):
            if iteration >= len(gen):
                break
            images_val, targets_val = batch[0], batch[1]

            with torch.no_grad():
                if cuda:
                    images_val  = images.type(torch.FloatTensor).cuda()
                    targets_val = torch.tensor(targets).cuda()
                else:
                    images_val  = images.type(torch.FloatTensor)
                    targets_val = torch.tensor(targets)

                images_val = F.interpolate(images_val, size=(3, 224, 224), mode='nearest')
                optimizer.zero_grad()
                outputs = nn.Sigmoid()(model_train(images_val))
                output  = loss(outputs, targets_val)

                _, preds = torch.max(outputs, 1)
                equal = torch.sum(preds.view(-1) == targets.view(-1))
                accuracy = equal / targets.__len__()

            val_loss            += output.item()
            val_total_accuracy  += accuracy.item()

            pbar.set_postfix(**{'val_loss'  : val_loss / (iteration + 1), 
                                'acc'       : val_total_accuracy / (iteration + 1)})#准确率
            pbar.update(1)
            
    print('Finish Validation')
    print('Epoch:'+ str(epoch+1) + '/' + str(Epoch))
    print('Total Loss: %.3f || Val Loss: %.3f ' % (total_loss / epoch_step, val_loss / epoch_step_val))
    #torch.save(model.state_dict()保存的是参数，torch.save(model)保存的是整个模型
    torch.save(model.state_dict(), 'logs/ep%03d-loss%.3f-val_loss%.3f.pth'%((epoch + 1), total_loss / epoch_step, val_loss / epoch_step_val))
